"""
Author : Mr.Sun
Datetime : 2024/11/6 10:28 
FileName : talk_2_db.py
Desc :  尝试使用mindsql 跟数据库进行对话
"""

# !pip install mindsql

from mindsql.core import MindSQLCore
from mindsql.databases import Sqlite
from mindsql.llms import GoogleGenAi
from mindsql.vectorstores import ChromaDB

# Add Your Configurations
config = {"api_key": "YOUR-API-KEY"}

# Choose the Vector Store. LLM and DB You Want to Work With And
# Create MindSQLCore Instance With Configured Llm, Vectorstore, And Database
minds = MindSQLCore(
    llm=GoogleGenAi(config=config),
    vectorstore=ChromaDB(),
    database=Sqlite()
)

# Create a Database Connection Using The Specified URL
connection = minds.database.create_connection(url="YOUR_DATABASE_CONNECTION_URL")

# Index All Data Definition Language (DDL) Statements in The Specified Database Into The Vectorstore
minds.index_all_ddls(connection=connection, db_name='NAME_OF_THE_DB')

# Index Question-Sql Pair in Bulk From the Specified Example Path
minds.index(bulk=True, path="your-qsn-sql-example.json")

# Ask a Question to The Database And Visualize The Result
response = minds.ask_db(
    question="YOUR_QUESTION",
    connection=connection,
    visualize=True
)

# Extract And Display The Chart From The Response
chart = response["chart"]
chart.show()

# Close The Connection to Your DB
connection.close()